package com.hyz.yuaiagent.rag.query;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.stereotype.Component;

/**
 * 使用@Component注解将此类标记为Spring框架中的一个组件，使其成为Spring管理的Bean。
 * 该类负责对查询进行重写，以提高信息检索的相关性。
 */
@Component
/**
 * 使用@Slf4j注解来自Lombok库，自动生成日志对象，方便在类中进行日志记录。
 */
@Slf4j
public class MyQueryRewriter {

    /**
     * ChatModel类型的私有成员变量，用于与聊天模型交互。
     */
    private final ChatModel dashscopeChatModel;

    /**
     * 构造方法，用于初始化MyQueryRewriter对象。
     *
     * @param dashscopeChatModel 一个ChatModel实例，用于与Dashscope聊天模型交互。
     */
    public MyQueryRewriter(ChatModel dashscopeChatModel) {
        this.dashscopeChatModel = dashscopeChatModel;
    }

    /**
     * 对输入的查询字符串进行重写。
     * 通过使用RewriteQueryTransformer来实现查询的重写，以期获得更相关的检索结果。
     *
     * @param query 需要重写的原始查询字符串。
     * @return 重写后的查询字符串。
     */
    public String doQueryRewrite(String query) {
        // 创建一个RewriteQueryTransformer实例，用于查询重写。
        RewriteQueryTransformer rewriteQueryTransformer = RewriteQueryTransformer.builder()
                .chatClientBuilder(ChatClient.builder(dashscopeChatModel))
                .build();

        // 使用transform方法对查询进行转换，并记录转换后的查询信息。
        Query transform = rewriteQueryTransformer.transform(new Query(query));
        log.info(transform.toString());

        // 返回转换后的查询文本。
        return transform.text();
    }
}
